/******************************************************************************

  File Name     : mxnet_model.h
  Version       : Initial Draft
  Author        : tongzhilin
  Created       : 2020/8/6
  Last Modified :
  Description   : mxnet_model
  Function List :
  History       :
  1.Date        : 2020/8/6
    Author      : tongzhilin
    Modification: Created file

******************************************************************************/
#pragma once
#include "mp_common.h"
#include "mxnet/c_predict_api.h"
#include "mxnet-cpp/MxNetCpp.h"

typedef struct
{
    INT32 m_s32N;      //max batch == cur batch == InputN in mxnet
    INT32 m_s32C;
    INT32 m_s32H;
    INT32 m_s32W;
}MXNET_SHAPE_SIZE_S;

typedef struct 
{
    char  m_cJsonFile[128];
    char  m_cParamsFile[128];
    INT32 m_s32DeviceType;              //The device type, 1: cpu, 2:gpu
    INT32 m_s32DeviceId;                //gpu id
    INT32 m_s32InputBlobNum;
    INT32 m_s32OutputBlobNum;
    MXNET_SHAPE_SIZE_S m_stInputShapeSize;
}MXNET_MODEL_INFO_S;

class MxnetModel
{
public:
    MxnetModel();
    virtual ~MxnetModel();

    static std::shared_ptr<MxnetModel> make_from_file(MXNET_MODEL_INFO_S* _pstModelInfo);
    static std::shared_ptr<MxnetModel> make_from_binary_stream();

    INT32 get_input_blob_num() const { return m_stModelInfo.m_s32InputBlobNum;}
    INT32 get_output_blob_num() const { return m_stModelInfo.m_s32OutputBlobNum;}

    PredictorHandle get_mxnet_model() 
    {
        return m_pMxnetHdl;
    }
    std::vector<UINT64> get_query_time()
    {
        return m_u64vQueryTimes;
    }

public:
    MXNET_MODEL_INFO_S      m_stModelInfo;
    //output blob
    std::vector<INT32>      m_s32vOutBlob_N;
    std::vector<INT32>      m_s32vOutBlob_C;
    std::vector<INT32>      m_s32vOutBlob_H;
    std::vector<INT32>      m_s32vOutBlob_W;
    std::vector<INT32>      m_s32vOutBlob_S;
private:
    void Init(char * _pcJsonBuffer, char *_pcParamsBuffer, INT32 _s32ParamsSize);
    void calculate_query_time();

    PredictorHandle         m_pMxnetHdl;        //mxnet model Predictor handle
    std::vector<UINT64>     m_u64vQueryTimes;
};
